from mopo.models.constructor import construct_single_model, load_env
import gym
import d4rl
import numpy as np
import tensorflow as tf
import mopo.mask
from shutil import copytree
import os

env_names = ["hopper-medium-replay-v0", "hopper-medium-v0", "hopper-random-v0"]

tf.reset_default_graph()


def stage_one(env_name, hidden_dim=50, model_save_dir="tmp"):
    tf.reset_default_graph()

    model_save_dirs = os.listdir(os.path.join("tmp", "model"))
    if env_name in model_save_dirs:
        return
    obs_dim, act_dim, inputs, outputs = load_env(env_name, use_diff_predict=False, normalize=True)

    model = construct_single_model(obs_dim=obs_dim, act_dim=act_dim,
                                   network_structure=np.ones([obs_dim + 1, obs_dim + act_dim]),
                                   hidden_dim=hidden_dim)
    model.train(inputs, outputs, batch_size=256, holdout_ratio=0.2)
    os.makedirs(model_save_dir, exist_ok=True)
    model.save(model_save_dir, timestep=0)


def stage_two(env_name, hidden_dim=50, model_save_dir="tmp"):
    tf.reset_default_graph()

    obs_dim, act_dim, inputs, outputs = load_env(env_name)

    model = construct_single_model(obs_dim=obs_dim, act_dim=act_dim, hidden_dim=hidden_dim, load_dir=model_save_dir,
                                   name="BNN_0")

    model.train(inputs, outputs, batch_size=256, holdout_ratio=0.2)
    mask = model.get_mask()
    return mask


def train(s1_env_name, s2_env_name, hidden_dim=100, save_dir="tmp"):
    os.makedirs(save_dir, exist_ok=True)
    os.makedirs(os.path.join(save_dir, "model"), exist_ok=True)
    os.makedirs(os.path.join(save_dir, "mask"), exist_ok=True)
    model_save_dir = os.path.join(save_dir, "model", "{}".format(s1_env_name))
    mask_save_dir = os.path.join(save_dir, "mask", "{}_{}.npy".format(s1_env_name, s2_env_name))
    print("**********\nstage one\n**********")
    stage_one(s1_env_name, hidden_dim, model_save_dir)
    print("**********\nstage two\n**********")
    mask = stage_two(s1_env_name, hidden_dim, model_save_dir)
    np.save(mask_save_dir, mask.astype(int))
    return mask


def mask_contrast(load_dir="tmp", base_env_name="hopper"):
    mask_path_list = os.listdir(os.path.join(load_dir, "mask"))
    mask_list = []
    for mask_path in mask_path_list:
        if mask_path.startswith(base_env_name):
            mask = np.load(os.path.join(load_dir, "mask", mask_path))
            # np.save(os.path.join(load_dir, "mask", mask_path), mask.astype(int))
            mask_list.append(mask)

    print((mask_list[0] == mask_list[1]).sum())

    complement_mask = np.zeros(mask_list[0].shape)
    for mask in mask_list:
        complement_mask[mask == 1] = 1
    np.save(os.path.join(load_dir, "mask", "{}_complement.npy".format(base_env_name)), complement_mask.astype(int))
    return complement_mask.astype(int)


if __name__ == "__main__":
    mask = train(env_names[2], env_names[0], hidden_dim=100)

    # complement_mask = mask_contrast()
    # print(complement_mask)
